{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9f8427b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-11-10T09:18:26.982999Z",
     "start_time": "2022-11-10T09:18:25.452775Z"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import torch\n",
    "import torchvision\n",
    "from torchvision import transforms, datasets\n",
    "from tqdm import tqdm\n",
    "import matplotlib.pyplot as plt\n",
    "from PIL import Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b211eee",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-11-10T09:18:26.989046Z",
     "start_time": "2022-11-10T09:18:26.985513Z"
    }
   },
   "outputs": [],
   "source": [
    "data_dir = '/data1/style_transfer/rec_split'\n",
    "train_dir = data_dir + '/train'\n",
    "test_dir = data_dir + '/test'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c00f6042",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-11-10T09:18:27.125897Z",
     "start_time": "2022-11-10T09:18:26.992347Z"
    }
   },
   "outputs": [],
   "source": [
    "# Batch size\n",
    "batch_size = 128\n",
    "# For faster computation, setting num_workers\n",
    "num_workers = 4\n",
    "\n",
    "# Transforms for the training, validation, and testing sets\n",
    "data_transforms = {\n",
    "    'train'      : transforms.Compose([transforms.RandomRotation(30),\n",
    "                                       transforms.RandomResizedCrop(224),\n",
    "                                       transforms.RandomHorizontalFlip(),\n",
    "                                       transforms.ToTensor(),\n",
    "                                       transforms.Normalize([0.485, 0.456, 0.406],\n",
    "                                                            [0.229, 0.224, 0.225])]),\n",
    "    'test'       : transforms.Compose([transforms.Resize(255),\n",
    "                                      transforms.CenterCrop(224),\n",
    "                                      transforms.ToTensor(),\n",
    "                                      transforms.Normalize([0.485, 0.456, 0.406],\n",
    "                                                           [0.229, 0.224, 0.225])])\n",
    "}\n",
    "\n",
    "# Loading the datasets with ImageFolder\n",
    "image_datasets = {\n",
    "    'train'  : datasets.ImageFolder(train_dir, transform=data_transforms['train']),\n",
    "    'test'   : datasets.ImageFolder(test_dir, transform=data_transforms['test'])\n",
    "}\n",
    "\n",
    "# Using the image datasets and the trainforms in defining the dataloaders\n",
    "dataloaders = {\n",
    "    'train' : torch.utils.data.DataLoader(image_datasets['train'], batch_size = batch_size, shuffle=True, num_workers = num_workers),\n",
    "    'test'  : torch.utils.data.DataLoader(image_datasets['test'], batch_size = batch_size)\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec007379",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-11-10T09:18:27.626319Z",
     "start_time": "2022-11-10T09:18:27.127787Z"
    }
   },
   "outputs": [],
   "source": [
    "model = torchvision.models.resnet50(pretrained=True)\n",
    "model.fc = torch.nn.Linear(in_features=2048, out_features=6)\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dbef448b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-11-10T09:18:27.719233Z",
     "start_time": "2022-11-10T09:18:27.629905Z"
    }
   },
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print('Is GPU available: ', 'Yes' if torch.cuda.is_available() else 'No')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7bc8e2a9",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-11-10T09:18:29.355576Z",
     "start_time": "2022-11-10T09:18:27.721957Z"
    }
   },
   "outputs": [],
   "source": [
    "criterion = torch.nn.CrossEntropyLoss()\n",
    "\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.0003)\n",
    "\n",
    "# Moving the model to the device\n",
    "model.to(device);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "539abe1e",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-11-10T09:21:14.284381Z",
     "start_time": "2022-11-10T09:18:29.358742Z"
    }
   },
   "outputs": [],
   "source": [
    "# Training the model\n",
    "epochs = 5\n",
    "steps = 0\n",
    "running_loss = 0\n",
    "print_every = 25\n",
    "for epoch in range(epochs):\n",
    "    qbar = tqdm(dataloaders['train'])\n",
    "    for inputs, labels in qbar:\n",
    "        steps += 1\n",
    "        # Move input and label tensors to the default device\n",
    "        inputs, labels = inputs.to(device), labels.to(device)\n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        logps = model.forward(inputs)\n",
    "        loss = criterion(logps, labels)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        running_loss += loss.item()\n",
    "        qbar.set_description('batch loss: %.3f'%loss.item())\n",
    "        \n",
    "    test_loss = 0\n",
    "    accuracy = 0\n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        for inputs, labels in dataloaders['test']:\n",
    "            inputs, labels = inputs.to(device), labels.to(device)\n",
    "            logps = model.forward(inputs)\n",
    "            batch_loss = criterion(logps, labels)\n",
    "            \n",
    "            test_loss += batch_loss.item()\n",
    "\n",
    "            # Calculate accuracy\n",
    "            ps = torch.exp(logps)\n",
    "            top_p, top_class = ps.topk(1, dim=1)\n",
    "            equals = top_class == labels.view(*top_class.shape)\n",
    "            accuracy += torch.mean(equals.type(torch.FloatTensor)).item()\n",
    "\n",
    "    print(f\"Epoch {epoch+1}/{epochs}.. \"\n",
    "          f\"Train loss: {running_loss/print_every:.3f}.. \"\n",
    "          f\"Validation loss: {test_loss/len(dataloaders['test']):.3f}.. \"\n",
    "          f\"Validation accuracy: { (accuracy/len(dataloaders['test']))*100 :.3f}%\")\n",
    "    running_loss = 0\n",
    "    model.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7204cb6c",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.13"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
